Add PointMass task and configurable optimizer for APG/PPO - #447
Add PointMass task and configurable optimizer for APG/PPO#447yangchen73 wants to merge 7 commits into
Conversation
There was a problem hiding this comment.
Pull request overview
Adds a new lightweight, differentiable RL “PointMass” task and extends the RL stack so APG (differentiable rollouts) and PPO/GRPO (standard rollouts) can share a unified train-rl entrypoint via rollout-kind routing. It also introduces configurable optimizers/LR schedulers, unifies evaluation logic across trainers, and adds scripts/configs/docs to support benchmarking and visualization.
Changes:
- Introduce
PointMassRLas a registered lightweight learning environment and provide APG/PPO training configs plus a trajectory visualization script. - Add optimizer + LR-scheduler configuration plumbing through
AlgorithmCfg/BaseAlgorithm, and update PPO/GRPO/APG to use the shared builders. - Unify deterministic evaluation and “best checkpoint” selection across
TrainerandDifferentiableTrainer, with routing based onRolloutKind.
Reviewed changes
Copilot reviewed 51 out of 51 changed files in this pull request and generated 1 comment.
Show a summary per file
| File | Description |
|---|---|
| tests/learning/test_routing.py | Adds coverage for routing trainer choice based on rollout_kind. |
| tests/learning/test_point_mass.py | Adds extensive unit tests for PointMass env behavior + unified training entry. |
| tests/learning/test_optimizer.py | Adds tests for optimizer and LR-scheduler builders/config coercion. |
| tests/learning/test_newton_planar_reach.py | Updates APG config usage to the new nested OptimizerCfg. |
| tests/learning/test_evaluation.py | Adds tests for unified deterministic evaluation semantics. |
| tests/learning/test_differentiable_trainer.py | Extends checkpoint coverage to include LR scheduler state; adds close() to mocks. |
| tests/learning/test_differentiable_env.py | Updates mock env protocol compliance with close(). |
| tests/learning/test_apg.py | Registers APG in algo registry, adds configurability + scheduler tests. |
| scripts/tutorials/rl/visualize_point_mass.py | Adds a plotting script for deterministic PointMass trajectories from checkpoints. |
| scripts/benchmark/rl/tasks/point_mass.yaml | Adds benchmark task definition for PointMass. |
| scripts/benchmark/rl/suites/point_mass.yaml | Adds benchmark suite for PointMass across APG/PPO. |
| scripts/benchmark/rl/runtime.py | Adds learning-env training path, unified evaluation usage, and routing integration. |
| scripts/benchmark/rl/algorithms/ppo.yaml | Migrates PPO benchmark config to nested optimizer config. |
| scripts/benchmark/rl/algorithms/grpo.yaml | Migrates GRPO benchmark config to nested optimizer config. |
| scripts/benchmark/rl/algorithms/apg.yaml | Adds APG benchmark algorithm config. |
| embodichain/learning/rl/utils/trainer.py | Unifies evaluation via evaluate_episodes, adds best-checkpoint selection and scheduler binding. |
| embodichain/learning/rl/utils/optimizer.py | Introduces shared optimizer/LR-scheduler builders and registries. |
| embodichain/learning/rl/utils/helper.py | Extends dict_to_tensordict to support tensor observations. |
| embodichain/learning/rl/utils/config.py | Adds OptimizerCfg and LRSchedulerCfg into the shared algorithm config. |
| embodichain/learning/rl/utils/init.py | Re-exports new optimizer/scheduler utilities and config objects. |
| embodichain/learning/rl/train.py | Adds unified learning-env training path and rollout-kind based trainer selection. |
| embodichain/learning/rl/routing.py | Adds get_trainer_class() to map rollout semantics to trainer type. |
| embodichain/learning/rl/models/init.py | Extends MLP builder options and orthogonal initialization support. |
| embodichain/learning/rl/experimental/newton/train_planar_reach.py | Migrates APG learning rate config to nested OptimizerCfg. |
| embodichain/learning/rl/evaluation.py | Adds shared deterministic asynchronous evaluation helper for both trainers. |
| embodichain/learning/rl/env.py | Adds learning-env registry (register_learning_env, build_learning_env) and shared protocol. |
| embodichain/learning/rl/differentiable_trainer.py | Adds eval/best-checkpoint parity with Trainer, plus scheduler checkpointing. |
| embodichain/learning/rl/algo/ppo.py | Switches PPO to shared optimization setup and scheduler stepping/logging. |
| embodichain/learning/rl/algo/grpo.py | Switches GRPO to shared optimization setup and scheduler stepping/logging. |
| embodichain/learning/rl/algo/base.py | Introduces RolloutKind + shared optimization/scheduler lifecycle hooks. |
| embodichain/learning/rl/algo/apg.py | Switches APG to shared optimization setup, adds rollout kind and scheduler stepping/logging. |
| embodichain/learning/rl/algo/init.py | Registers APG and normalizes nested optimizer/scheduler config mappings. |
| embodichain/learning/rl/init.py | Re-exports new top-level RL APIs (learning env, evaluation, routing). |
| embodichain_tasks/embodichain_tasks/rl/basic/point_mass.py | Adds the differentiable PointMass environment and registers it as PointMassRL. |
| embodichain_tasks/embodichain_tasks/rl/basic/init.py | Exports PointMass from the basic RL tasks package. |
| embodichain_tasks/configs/agents/rl/push_cube/train_config.json | Migrates PPO config to nested optimizer block. |
| embodichain_tasks/configs/agents/rl/push_cube/train_config_grpo.json | Migrates GRPO config to nested optimizer block. |
| embodichain_tasks/configs/agents/rl/basic/point_mass/train_ppo.yaml | Adds PPO training config for PointMass learning env. |
| embodichain_tasks/configs/agents/rl/basic/point_mass/train_apg.yaml | Adds APG training config for PointMass learning env. |
| embodichain_tasks/configs/agents/rl/basic/cart_pole/train_config.yaml | Migrates PPO config to nested optimizer block. |
| embodichain_tasks/configs/agents/rl/basic/cart_pole/train_config.json | Migrates PPO config to nested optimizer block. |
| embodichain_tasks/configs/agents/rl/basic/cart_pole/train_config_grpo.yaml | Migrates GRPO config to nested optimizer block. |
| embodichain_tasks/configs/agents/rl/basic/cart_pole/train_config_grpo.json | Migrates GRPO config to nested optimizer block. |
| docs/source/tutorial/rl.rst | Documents trainer.learning_env vs trainer.gym_config and adds PointMass examples. |
| docs/source/overview/rl/trainer.md | Updates Trainer docs and adds DifferentiableTrainer parity notes. |
| docs/source/overview/rl/train_script.md | Documents unified routing and PointMass usage in train.py. |
| docs/source/api_reference/embodichain/embodichain.learning.rl.rst | Adds API reference sections for env/eval/routing + trainer. |
| docs/source/api_reference/embodichain/embodichain.learning.rl.algo.rst | Updates algo overview to include APG and RolloutKind. |
| agent_context/topics/rl-learning/rl-learning.md | Updates agent context docs for learning env registry, evaluation, and APG routing. |
| agent_context/MAP.yaml | Updates “rl-learning” topic sources to include routing/evaluation/PointMass. |
| .agents/skills/add-task-env/SKILL.md | Updates skill docs to reflect embodichain_tasks layout and learning-env tasks. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| metrics = info.get("metrics", {}) | ||
| if isinstance(metrics, Mapping): | ||
| for name, value in metrics.items(): | ||
| metric_values.setdefault(str(name), []).extend( | ||
| _selected_values(value, selected) | ||
| ) | ||
|
|
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 52 out of 52 changed files in this pull request and generated no new comments.
Suppressed comments (4)
scripts/benchmark/rl/runtime.py:490
- evaluate_checkpoint() binds an evaluation rollout buffer via env.set_rollout_buffer(), but the new evaluate_episodes() loop does not populate policy-side fields (notably previous actions) into that buffer. This can corrupt rewards/metrics for environments whose managers read prev actions from env.rollout_buffer (e.g. action_smoothness_penalty reads buffer["action"]). Prefer not binding a rollout buffer in this evaluation path unless the evaluator also writes the buffer fields each step.
scripts/benchmark/rl/runtime.py:505 - The on_step callback currently only compacts eval_rollout_buffer, but never writes action/reward/done into the buffer. After removing rollout-buffer binding for evaluation, this callback can be dropped and evaluate_episodes() called without on_step.
embodichain/learning/rl/train.py:98 - _build_learning_policy() computes policy_name but never uses it. This is dead code and can trigger lint/type-check warnings; remove it (or use it for validation) to keep the entry point clean.
obs_dim = int(env.single_observation_space.shape[-1])
action_dim = int(env.single_action_space.shape[-1])
policy_name = policy_block["name"].lower()
actor_cfg = policy_block.get("actor")
critic_cfg = policy_block.get("critic")
embodichain_tasks/embodichain_tasks/rl/basic/point_mass.py:33
- PointMassEnv is newly introduced, but it registers with override=True. This silently replaces any existing registration for "PointMassRL" and can mask accidental duplicate IDs during imports. Prefer the default (override=False) and rely on the registry's nested-package shadow detection for editable installs.
| episode metrics. | ||
| - `save_checkpoint()`: Save model parameters and training state. | ||
|
|
||
| ### DifferentiableTrainer |
There was a problem hiding this comment.
We should add docs for APG at algorithm section
| @@ -0,0 +1,328 @@ | |||
| # ---------------------------------------------------------------------------- | |||
There was a problem hiding this comment.
This script should not be included in official tutorial.
|
Check that can DDP training work correctly for APG |
|
Also check whether the profiler can be correctly used for RL |
Description
Type of change
Screenshots
Checklist
black .command to format the code base.